# -*- coding: utf-8 -*-
"""
AI写作服务
"""

import os
from typing import List, Optional

try:
    from langchain_openai import ChatOpenAI
    LANGCHAIN_OPENAI_AVAILABLE = True
except ImportError:
    LANGCHAIN_OPENAI_AVAILABLE = False

from ..models.character import Character


class AIWriter:
    """AI写作服务类"""
    
    def __init__(self, api_key: str = None, base_url: str = None, model: str = None, 
                 temperature: float = None, max_tokens: int = None):
        """初始化AI写作服务
        
        Args:
            api_key: OpenAI API密钥
            base_url: API基础URL，用于使用其他兼容的API服务
            model: AI模型名称
            temperature: 创作温度（0-2）
            max_tokens: 最大token数
        """
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        self.base_url = base_url or os.getenv("OPENAI_BASE_URL")
        self.model = model or "gpt-3.5-turbo"
        self.temperature = temperature if temperature is not None else 0.8
        self.max_tokens = max_tokens
        self.client = None
        self.available = False
        
        if self.api_key:
            self._init_client()
    
    def _init_client(self):
        """初始化客户端"""
        if not LANGCHAIN_OPENAI_AVAILABLE:
            self.available = False
            return
        
        try:
            # 构建LangChain OpenAI参数
            kwargs = {
                "api_key": self.api_key,
                "model": self.model,
                "temperature": self.temperature,
            }
            
            # 如果有自定义base_url，添加到参数中
            if self.base_url:
                kwargs["base_url"] = self.base_url
            
            # 如果设置了最大token数
            if self.max_tokens:
                kwargs["max_tokens"] = self.max_tokens
            
            self.client = ChatOpenAI(**kwargs)
            self.available = True
        except Exception as e:
            print(f"AI服务初始化失败: {e}")
            self.available = False
    
    def is_available(self) -> bool:
        """检查AI服务是否可用"""
        return self.available and LANGCHAIN_OPENAI_AVAILABLE
    
    def generate_story(self, 
                      plot: str, 
                      characters: List[Character], 
                      target_word_count: int = 1000,
                      style: str = "小说") -> Optional[str]:
        """生成故事内容
        
        Args:
            plot: 大致剧情
            characters: 参与的角色列表
            target_word_count: 目标字数
            style: 写作风格
            
        Returns:
            生成的故事内容，失败返回None
        """
        if not self.is_available():
            return None
        
        try:
            # 构建角色信息
            character_info = []
            for char in characters:
                info = f"角色名：{char.name}"
                if char.description:
                    info += f"\n描述：{char.description}"
                
                # 添加属性信息
                attrs_text = char.get_attributes_text()
                if attrs_text:
                    info += f"\n属性：\n{attrs_text}"
                
                character_info.append(info)
            
            characters_text = "\n\n".join(character_info) if character_info else "无指定角色"
            
            # 计算字数范围
            min_words = int(target_word_count * 0.9)  # 最少90%
            max_words = int(target_word_count * 1.1)  # 最多110%
            
            # 构建更精确的提示信息
            prompt = f"""你是一位专业的{style}作家，现在需要为小说的一个章节创作内容。

【重要要求】：
1. 严格控制字数在{min_words}-{max_words}字之间，目标字数为{target_word_count}字
2. 必须充分运用每个角色的详细属性，让角色的性格、技能、职业等特点在故事中自然体现
3. 对话要完全符合角色的性格特点和职业背景
4. 情节要紧扣剧情大纲，推动故事发展
5. 文笔要流畅，符合{style}的文体特点
6. 直接输出小说正文内容，不要任何说明文字、标题或总结

【剧情大纲】：
{plot}

【角色信息】：
{characters_text}

【创作要求】：
- 这是小说章节的一部分内容，需要有完整的场景描写
- 角色行为和对话必须体现其属性特点
- 情节要有起承转合，不能只是简单描述
- 字数务必达到{target_word_count}字左右

请开始创作章节内容："""

            # 调用LangChain AI生成
            response = self.client.invoke(prompt)
            
            # 返回生成的内容，去除可能的说明文字
            content = response.content.strip()
            
            # 简单清理可能的说明文字（开头的标题或结尾的总结）
            lines = content.split('\n')
            clean_lines = []
            
            for line in lines:
                line = line.strip()
                # 过滤掉明显的说明性文字
                if (line.startswith('【') or line.startswith('【') or 
                    line.startswith('注：') or line.startswith('说明：') or
                    line.startswith('总结：') or line.startswith('以上') or
                    '字数统计' in line or '创作完成' in line):
                    continue
                if line:
                    clean_lines.append(line)
            
            return '\n'.join(clean_lines)
            
        except Exception as e:
            print(f"AI生成故事失败: {e}")
            return None
    
    def improve_text(self, text: str, instruction: str = "润色以下文本") -> Optional[str]:
        """改进文本
        
        Args:
            text: 原始文本
            instruction: 改进指令
            
        Returns:
            改进后的文本，失败返回None
        """
        if not self.is_available():
            return None
        
        try:
            # 构建提示
            if text:
                prompt = f"{instruction}：\n\n{text}"
            else:
                # 如果没有文本，指令本身就是prompt（用于AI名字生成等）
                prompt = instruction
            
            response = self.client.invoke(prompt)
            
            return response.content.strip()
            
        except Exception as e:
            print(f"AI改进文本失败: {e}")
            return None
    
    def set_api_config(self, api_key: str, base_url: str = None, model: str = None,
                      temperature: float = None, max_tokens: int = None):
        """设置API配置"""
        self.api_key = api_key
        self.base_url = base_url
        self.model = model or "gpt-3.5-turbo"
        self.temperature = temperature if temperature is not None else 0.8
        self.max_tokens = max_tokens
        
        if not LANGCHAIN_OPENAI_AVAILABLE:
            print("LangChain OpenAI库未安装")
            self.available = False
            return False
        
        # 重新初始化客户端
        self._init_client()
        return self.available
    
    def get_config_info(self) -> dict:
        """获取当前配置信息"""
        return {
            "model": self.model,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "base_url": self.base_url,
            "available": self.is_available()
        } 